package com.tianji.aigc.service.impl;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.tianji.aigc.config.SystemPromptConfig;
import com.tianji.aigc.config.ToolResultHolder;
import com.tianji.aigc.constants.Constant;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.service.ChatService;
import com.tianji.aigc.service.ChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.utils.UserContext;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Service
@RequiredArgsConstructor
public class ChatServiceImpl implements ChatService {

    private final ChatClient chatClient;
    private final ChatClient titleClient;
    private final SystemPromptConfig systemPromptConfig;
    private final ChatMemory redisChatMemory;

    private final VectorStore vectorStore;

    private final ChatSessionService chatSessionService;

    // 作为一个全局的、线程安全的状态存储，用来标记特定 sessionId 对应的对话生成任务是否应该继续进行
    // Key: sessionId，唯一标识一个用户的会话
    // Value: Boolean，true 表示生成任务应该继续，false 或 null (通过 getOrDefault) 表示生成任务应该被停止
    // 目前的版本暂时用Map实现，如果考虑分布式环境的话，可以考虑用redis来实现
    private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();

    private static final ChatEventVO STOP_EVENT = ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build();

     /**
     * 执行聊天对话功能，接收用户问题并返回流式的聊天响应
     *
     * @param question 用户提出的问题
     * @param sessionId 会话标识符，用于标识当前对话上下文
     * @return 返回包含聊天事件的流，包括数据事件（聊天内容）和结束事件（如停止事件）
     */
    @Override
    public Flux<ChatEventVO> chat(String question, String sessionId) {
        // 获取当前会话对应的 conversationId，用于后续对话上下文管理
        String conversationId = ChatService.getConversationId(sessionId);

        // 用于累积输出内容，便于在取消时保存历史记录
        StringBuilder output = new StringBuilder();

        // 生成本次请求的唯一标识，用于工具调用上下文关联
        String requestId = IdUtil.fastSimpleUUID();

        // 获取用户id
        Long userId = UserContext.getUser();

        // 异步更新会话信息
        //this.chatSessionService.update(sessionId, question, userId);

        // 构建问答顾问对象，配置向量检索参数：相似度阈值为0.6，最多返回6个相关结果
        var qaAdvisor = QuestionAnswerAdvisor.builder(this.vectorStore)
                .searchRequest(SearchRequest.builder().similarityThreshold(0.6d).topK(6).build())
                .build();

        // 构建聊天请求，设置用户问题、系统提示信息以及上下文顾问
        return chatClient.prompt()
                .toolContext(Map.of(Constant.REQUEST_ID, requestId, Constant.USER_ID, userId))
                .user(question)
                .system(
                        promptSystemSpec -> promptSystemSpec
                            .text(systemPromptConfig.getChatSystemPrompt().get())
                            .param("now", DateUtil.now())
                )
                .advisors(advisor -> advisor
                        .advisors(qaAdvisor)
                        .param(ChatMemory.CONVERSATION_ID, conversationId))
                // 执行流式聊天请求并处理响应
                .stream()
                .chatResponse()
                // 在流开始前设置会话生成状态为 true，表示正在生成中
                .doFirst(() -> {GENERATE_STATUS.put(sessionId, true);})
                // 发生错误时清除会话生成状态，避免状态残留
                .doOnError(throwable -> {GENERATE_STATUS.remove(sessionId);})
                // 流完成时清除会话生成状态，表示生成结束
                .doOnComplete(() -> {GENERATE_STATUS.remove(sessionId);})
                // 在取消时保存当前已输出的内容作为历史记录
                .doOnCancel(() -> saveStopHistoryRecord(conversationId, output))
                // 根据会话状态控制流的持续性，若状态为 false 则终止流
                .takeWhile(response -> GENERATE_STATUS.getOrDefault(sessionId, false))
                .doFinally(signalType -> {
                    this.chatSessionService.updateByAI(sessionId, output.toString(), question, userId);
                })
                // 将聊天响应转换为数据事件，逐段返回聊天内容
                .map(response -> {
                    // 处理响应结果：如果是最后一条数据，则将消息ID与requestId进行绑定，
                    // 以便后续可以通过消息ID获取请求参数，解决历史记录中缺失参数的问题
                    String finishReason = response.getResult().getMetadata().getFinishReason();
                    if (StrUtil.equals(Constant.STOP, finishReason)) {
                        String messageId = response.getMetadata().getId();
                        ToolResultHolder.put(messageId, Constant.REQUEST_ID, requestId);
                    }

                    String text = response.getResult().getOutput().getText();
                    output.append(text);
                    return ChatEventVO.builder()
                            .eventData(text)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                })
                // 在数据流结束后添加停止事件，并处理可能存在的工具调用参数返回
                .concatWith(
                        Flux.defer( () -> {
                            // 通过请求id获取参数列表，如果不为空则将其封装成PARAM事件追加到结果流中
                            var map = ToolResultHolder.get(requestId);
                            if (MapUtil.isNotEmpty(map)) {
                                ChatEventVO chatEventVO = ChatEventVO.builder()
                                        .eventData(map)
                                        .eventType(ChatEventTypeEnum.PARAM.getValue())
                                        .build();
                                ToolResultHolder.remove(requestId);
                                return Flux.just(chatEventVO,STOP_EVENT);
                            }
                            return Flux.just(STOP_EVENT);
                        })
                );
    }




    /**
     * 保存停止输出的记录
     *
     * @param conversationId 会话id
     * @param output        大模型输出的内容
     */
    private void saveStopHistoryRecord(String conversationId, StringBuilder output) {
        redisChatMemory.add(conversationId, new AssistantMessage(output.toString()));
    }


    /**
     * 停止当前会话的聊天
     *
     * @param sessionId 会话id
     */
    @Override
    public void stopChat(String sessionId) {
        GENERATE_STATUS.remove(sessionId);
    }
}
